package com.xiashitech.agent.instrumentation.biz.userinfo.httprequest.inst;

import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.biz.userinfo.httprequest.UserInfoHelper;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.httprequest.XSHttpServletRequestWrapper;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import com.xiashitech.agent.utils.DurationPrinter;
import com.xiashitech.agent.utils.XSUtil;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.springframework.web.context.request.ServletWebRequest;

import javax.servlet.http.HttpServletRequest;

import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSSpringInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("org.springframework.web.servlet.mvc.method.annotation.ServletInvocableHandlerMethod");
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                isMethod()
                        .and(isPublic())
                        .and(named("invokeAndHandle"))
                        .and(takesArgument(0, named("org.springframework.web.context.request.ServletWebRequest"))),
                XSSpringInstrumentation.class.getName() + "$InnerAdvice");
    }

    @SuppressWarnings("unused")
    public static class InnerAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Argument(value = 0) ServletWebRequest servletWebRequest,
                @Advice.Local("user") UserRequest userRequest,
                @Advice.Local("otelScope") Scope scope) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSpringInstrumentation");
            HttpServletRequest httpServletRequest = servletWebRequest.getRequest();
            if (!(httpServletRequest instanceof XSHttpServletRequestWrapper)){
                return;
            }
            if (XSUtil.isEmpty(Java8Bridge.currentForBaggage().getEntryValue("xs.biz.state"))){
                return;
            }
            userRequest = UserRequest.builder().request(httpServletRequest).build();
            UserInfoHelper.captureUserId(userRequest);
            scope = BaggageUtil.syncBaggage().build().makeCurrent();
            durationPrinter.printElapseDuration("InnerAdvice b");
        }
    }

    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
    public static void onExit(
            @Advice.Local("otelScope") Scope scope,
            @Advice.Local("user") UserRequest userRequest) {
        DurationPrinter durationPrinter = new DurationPrinter("XSSpringInstrumentation");
        try {
            if (scope == null) {
                return;
            }
            scope.close();
        } catch (Throwable e) {
            StartConfig.throwableStackPrinter(e);
        }
        durationPrinter.printElapseDuration("InnerAdvice e");
    }

}
